from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain_openai import ChatOpenAI
from pydantic import SecretStr
from langchain_core.output_parsers import StrOutputParser
from langchain_base_prompt import Base_prompt
from langchain.chains import create_retrieval_chain
from langchain_RAG_test import RAGSystem
from langchain.schema import HumanMessage, AIMessage
from typing import List


class Main:
    def __init__(self):
        self.prompt = Base_prompt().get_base_prompt()
        self.llm = ChatOpenAI(
            base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
            api_key=SecretStr("sk-1602eeff7d0442f5b814aaf051af33fc"),
            model="qwen-max",
            streaming=False, # 流式输出
            temperature=0.7 # 温度
        )

        # 创建链
        self.document_chain = create_stuff_documents_chain(self.llm, self.prompt)
        self.vector = RAGSystem.get_vector_store()
        self.retriever = self.vector.as_retriever()
        self.retrieval_chain = create_retrieval_chain(self.retriever, self.document_chain)
        
        # 对话历史
        self.conversation_history: List[dict] = []
    

    # 获取检索链
    def get_retrieval_chain(self):
        return self.retrieval_chain
    
    # 对话
    def chat(self, user_input: str) -> dict:
        """
        进行对话，保持历史记录
        """
        # 获取检索结果
        response = self.retrieval_chain.invoke({"input": user_input})
        
        # 记录对话历史
        self.conversation_history.append({
            "user": user_input,
            "context": response['context'],
            "ai_response": response['answer'],
            "timestamp": response.get('timestamp', '')
        })
        
        return response
    
    def get_conversation_history(self) -> List[dict]:
        """
        获取对话历史
        """
        return self.conversation_history
    
    def clear_history(self):
        """
        清空对话历史
        """
        self.conversation_history = []
    
    def print_conversation_summary(self):
        """
        打印对话摘要
        """
        print(f"\n=== 对话历史摘要 ===")
        print(f"总对话轮数: {len(self.conversation_history)}")
        if self.conversation_history:
            print(f"最近问题: {self.conversation_history[-1]['user']}")
        print("=" * 30)

def interactive_chat():
    """
    交互式对话函数
    """
    main = Main()
    
    print("=== RAG智能对话系统 ===")
    print("输入 'quit' 退出，输入 'history' 查看历史，输入 'clear' 清空历史")
    print("=" * 50)
    
    while True:
        try:
            # 获取用户输入
            user_input = input("\n你: ").strip()
            
            # 处理特殊命令
            if user_input.lower() == 'quit':
                print("再见！")
                break
            elif user_input.lower() == 'history':
                # 显示对话历史
                history = main.get_conversation_history()
                print("\n=== 对话历史 ===")
                for i, conv in enumerate(history[-5:], 1):  # 显示最近5轮对话
                    print(f"\n第{i}轮:")
                    print(f"你: {conv['user']}")
                    print(f"AI: {conv['ai_response'][:100]}...")
                print("=" * 30)
                continue
            elif user_input.lower() == 'clear':
                # 清空对话历史
                main.clear_history()
                print("对话历史已清空")
                continue
            elif not user_input:
                continue
            
            # 进行对话
            print("AI正在思考...")
            response = main.chat(user_input)
            
            # 显示AI回复
            print(f"AI: {response['answer']}")
            
            # 显示对话摘要
            main.print_conversation_summary()
            
        except KeyboardInterrupt:
            print("\n\n再见！")
            break
        except Exception as e:
            print(f"发生错误: {e}")

if __name__ == "__main__":
    # 运行交互式对话
    interactive_chat()
    
    # 或者运行单次对话（原来的方式）
    # main = Main()
    # retrieval_chain = main.get_retrieval_chain()
    # response = retrieval_chain.invoke({"input": "您好，请问您是"})
    # print(response['answer'])